DeepSeek FP8 Training

DeepSeek V3 implemented FP8 GEMM operator for training, i.e., keep params in FP8 during training. However, given that some operators require full precision to keep stability during computation, DeepSeek V3 maintains full precision or half precision (FP32 or BF16) on several operators: Embedding layer, Output Head, MoE Gating, Normalization, Attention Module. Only linear layers uses FP8 GEMM.

The overall training adopts mixed precision schema. Besides FP8 quantized weights, master weights, weight gradients and optimizer states are still stored with higher precision, and are dynamically casted to lower precision for computation.

For example, the above illustrates the FP8 computation pipeline:

1. Activition from previous layer in BF16 format is quantized to FP8
2. Master weight (FP32) is quantized to FP8. Then the quantized version is involved in FP8 GEMM.
3. FP8 GEMM computes intermediate results (tile-based computation), and accumulates them with FP32 register for accuracy concerns.
4. Then, the result in FP32 is quantized to BF16 and stored to output matrix.

During backward propagation, to update the weight:

1. The gradient of output from next layer (BF16) is quantized to FP8. Also, the activition matrix (input) is also quantized to FP8 and sent to Wgrad FP8 GEMM. We also use FP32 accumulator due to precision concerns.
2. The output is Weight Gradient. We store it in FP32. Later, the gradient is converted to BF16 to update the optimizer states.
3. The optimizer states are transformed to FP32 to update master weight.

On the other side, to get the gradient of input:

1. The master weight is converted to FP8. The output gradient is also converted to FP8. These 2 matrices are passed to FP8 GEMM operator for Dgrad. We also use FP32 for accumulator.
2. The result is quantized to BF16 to obtain Input Gradient. The backprop continues.

The challenge of implementing FP8 GEMM kernel includes

• underflow and overflow
• presence of outliers in activition, weights and gradients

So DeepSeek’s solution for this is fine-grained quantization that quantize values within smaller groups.

• tile-wise \(1\times N_c\) grouping for activition
• block-wise \(N_c \times N_c\) grouping for weights

Why different grouping size for activition and weight? The paper says the reason is that Dgrad is sensitive to precision, \(N_c\times N_c\) blockwise quantization will make model diverge on MoE. In addition, “activition gradients are highly imbalanced among tokens, resulting in token-correlated outliers”, as suggested by paper.